// 导入可视化库，用来画图的
import * as tfvis from '@tensorflow/tfjs-vis';
// 导入TensorFlow.js，核心库
import * as tf from '@tensorflow/tfjs';

// 页面加载完了就开始搞事情
window.onload = async () => {
  // 随便弄点简单的数据，x是输入，y是输出，一看就知道y=2x-1
  const xs = [1,2,3,4]
  const ys = [1,3,5,7]
  // 打印一下可视化库，看看是不是正常加载了
  console.log(tfvis)
  
  // 画个散点图看看数据长啥样
  tfvis.render.scatterplot(
    {name: '线性回归训练集'},
    {values: xs.map((x, i) => ({x, y: ys[i] }))},
    {xAxisDomain: [0, 5], yAxisDomain: [0, 8]} // 设置坐标轴范围，看起来舒服点
  )
  
  // 创建一个顺序模型，就是一层一层堆上去的那种
  const model = tf.sequential()
  // 添加一个全连接层，1个输出单元，输入是1维的
  model.add(tf.layers.dense({units: 1, inputShape: [1]}))
  
  // 编译模型，用均方误差当损失函数，随机梯度下降当优化器，学习率设为0.1（随便设的）
  model.compile({loss: tf.losses.meanSquaredError, optimizer: tf.train.sgd(0.1)}) 
  
  // 把JS数组转成TensorFlow的张量，模型只认这个格式
  const inputs = tf.tensor(xs)
  const labels = tf.tensor(ys)
  
  // 开始训练啦，一次喂4个数据，训100轮，顺便用回调函数把训练过程画出来看看
  await model.fit(inputs, labels, {
    batchSize: 4,
    epochs: 100,
    callbacks: tfvis.show.fitCallbacks(
      {name: '训练过程'},
      ['loss'], // 只看损失函数的变化就行
    )
  })
  const output = model.predict(tf.tensor([5]))
  console.log(output.dataSync()[0])
}